from BitArray import BitArray

class leaf_class:
    def __init__(self, obj):
        self._symb = obj[0]
        self._weight = obj[1]

    def __str__(self):
        return str(self._symb)+','+str(self._weight)
    
    def getSymbols(self):
        return self._symb

    def getWeight(self):
        return self._weight

class tree_class:
    def __init__(self, left, right):
        self._right = right
        self._left = left
        self._symb = left.getSymbols()+right.getSymbols()
        self._weight = left.getWeight()+right.getWeight()
        
    def left_branch(self):
        return self._left

    def right_branch(self):
        return self._right

    def getWeight(self):
        return self._weight

    def getSymbols(self):
        return self._symb

def insort(lst, x):
    lo = 0
    hi = len(lst)
    while lo < hi:
        mid = (lo+hi)/2
        if lst[mid].getWeight() > x.getWeight():
            lo = mid+1
        else:
            hi = mid
    lst.insert(lo,x)

def make_tree(obj):
    lst = []
    for x in obj:
        insort(lst, leaf_class(x))
    while len(lst) > 1:
        insort(lst, tree_class(lst.pop(), lst.pop()))
    return lst[0]

def frequency_table(msg):
    table = []
    i = 0
    while i < 256:
        x = chr(i)
        freq = 1
        for ch in msg:
            if ch == x: freq = freq + 1
        table.append((x, freq))
        i = i + 1
    return table

def decode_byte(bits, t):
    if isinstance(t, leaf_class):
        return t.getSymbols()
    if len(bits) == 0:
        return None
    if bits.front() == 1:
        bits.pop_front()
        return decode_byte(bits,t.right_branch())
    elif bits.front() == 0:
        bits.pop_front()
        return decode_byte(bits,t.left_branch())
    else: raise ValueError('Symbol can not be decoded')

def encode_byte(symb, bitarray, t):
    while isinstance(t, tree_class):
        if symb in t.left_branch().getSymbols():
            t = t.left_branch()
            bitarray.push_back(0)
        elif symb in t.right_branch().getSymbols():
            t = t.right_branch()
            bitarray.push_back(1)
        else: raise ValueError('Symbol can not be encoded')
    return bitarray

def encode_file(inputname, outputname):
    fin = open(inputname,'rb')
    msg = fin.read()
    table = frequency_table(msg)
    tree = make_tree(table)
    ba = BitArray()
    for x in msg:
        encode_byte(x, ba, tree)
    fin.close()
    fout = open(outputname,'w')
    fout.write(str(table)+'\n')
    fout.write(ba.toString())
    fout.close()
    
def decode_file(inputname, outputname):
    fin = open(inputname, 'rb')
    s = "table = " + fin.readline()
    exec(s)
    front = int(fin.readline())
    back = int(fin.readline())
    s = fin.read()
    ba = BitArray()
    ba.fromString(front, back, s)
    tree = make_tree(table)
    fin.close()
    fout = open(outputname, 'wb')
    while len(ba)>0:
        fout.write(decode_byte(ba, tree))
    fout.close()
